from env.mec_env_v2 import MECEnv
import wandb
import opt_algo.baseline_mig as baseline
import opt_algo.MultiDqn as multiDQN
import opt_algo.mec_sac_discrete as SACDM
import threading



def always(env, user_num, server_num):
    env.reset()
    wandb.init(
        project="ServiceMigration",
        name=f"ALWAYS_{user_num}_{server_num}"
    )
    baseline.train(env, baseline.always_policy)
    wandb.finish()


def nomig(env, user_num, server_num):
    env.reset()
    wandb.init(
        project="ServiceMigration",
        name=f"NOMIG_{user_num}_{server_num}"
    )
    baseline.train(env, baseline.no_mig_policy)
    wandb.finish()


def random(env, user_num, server_num):
    env.reset()
    wandb.init(
        project="ServiceMigration",
        name=f"RANDOM_{user_num}_{server_num}"
    )
    baseline.train(env, baseline.random_policy)
    wandb.finish()

def dqn(env, user_num, server_num):
    env.reset()
    wandb.init(
        project="ServiceMigration",
        name=f"DQN_{user_num}_{server_num}"
    )
    multiDQN.run(user_num, server_num)
    wandb.finish()

def sac(env, user_num, server_num):
    env.reset()
    wandb.init(
        project="ServiceMigration",
        name=f"SACDM_{user_num}_{server_num}"
    )
    SACDM.run(user_num, server_num)
    wandb.finish()

def exp(user_num, server_num):
    env = MECEnv.env(user_num, server_num)
    threading.Thread(target=always, args=(env, user_num, server_num)).start()
    threading.Thread(target=nomig, args=(env, user_num, server_num)).start()
    threading.Thread(target=random, args=(env, user_num, server_num)).start()
    # threading.Thread(target=dqn, args=(env, user_num, server_num)).start()
    # threading.Thread(target=sac, args=(env, user_num, server_num)).start()




if __name__ == '__main__':
    user_num = 6
    server_num = 6
    env = MECEnv.env(user_num, server_num)
    # always(env, user_num, server_num)
    # random(env, user_num, server_num)
    # nomig(env, user_num, server_num)
    # dqn(env,user_num, server_num)
    sac(env, user_num, server_num)
